Implementing a neural network interatomic model with performance portability for emerging exascale architectures

The two main thrusts of computational science are increasingly accurate predictions and faster calculations; to this end, the zeitgeist in molecular dynamics (MD) simulations is pursuing machine learned data driven interatomic models, e.g. neural network potentials, novel hardware architectures, GPUs. Current implementations potentials orders magnitude slower than traditional models while looming exascale computing offers ability run large, with these achieving portable performance for MD new varied requires rethinking algorithms, using structures, library solutions. We re-implement a model CabanaMD, an proxy application, built on libraries developed portability. Our implementation shows significantly improved thread scaling complex kernel as compared current LAMMPS implementation, across both strong weak scaling. single-source solution enables up 20 million atoms single CPU node 4 GPU. also explore parallelism layout choices (using flexible structures called AoSoAs) their effect performance, seeing ∼50% ∼5% improvements GPU by choosing right level respectively. Program title: CabanaMD-NNP CPC Library link program files: https://doi.org/10.17632/x948kyy7jh.1 Developer's repository link: https://github.com/ECP-CoPA/CabanaMD, https://github.com/CompPhysVienna/n2p2 Licensing provisions: BSD3-Clause, GPL-3.0 Programming Language: C++ Nature problem: Developing potential architectures. Solution method: uses algorithms data-structures from Kokkos [1] Cabana [2] computations Behler-Parrinello [3, 4] portability hardware. All stored atomic properties array-of-structs-of-arrays (Cabana::AoSoAs), auxiliary values including parameters arrays (Kokkos::Views). computation done way: propagation parallel kernels (Kokkos::parallel_for), calculations performed each atom neighbor, evaluation descriptors (symmetry functions) forces, use extensions constructs (Cabana::neighbor_parallel_for). These provide our significant speedups CPUs GPUs large systems, additionally allowing flexibility further optimizations. Additional comments restrictions unusual features: previously n2p2 package [4] contains interface [5], which we compare throughout paper. primarily extend directly (https://github.com/CompPhysVienna/n2p2) add that extension within CabanaMD (https://github.com/ECP-CoPA/CabanaMD), obtain results identical input file. H.C. Edwards, C.R. Trott, D. Sunderland, Kokkos: Enabling manycore through polymorphic memory access patterns, Journal Parallel Distributed Computing. 74 (2014) 3202–3216. https://doi.org/10.1016/j.jpdc.2014.07.003. S. Slattery, C. Junghans, Lebrun-Grandie, R. Halver, G. Chen, Reeve, A. Scheinberg, Smith, Bird, ECP-CoPA/Cabana: Version 0.2.0, Zenodo, 2019. J. Behler, M. Parrinello, Generalized Neural-Network Representation High-Dimensional Potential-Energy Surfaces, Phys. Rev. Lett. 98 (2007) 146401. Singraber, Dellago, Library-Based Implementation Neural Network Potentials, Chem. Theory Comput. 15 (2019) 1827–1840. https://doi.org/10.1021/acs.jctc.8b00770. Plimpton, Fast Algorithms Short-Range Molecular Dynamics, Computational Physics. 117 (1995) 1–19. https://doi.org/10.1006/jcph.1995.1039.